from typing import List

from langchain.chains import LLMChain
from langchain.output_parsers import PydanticOutputParser, OutputFixingParser
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

from backend.app.schemas.device import DeviceAttributeListSchema
from backend.app.schemas.search import OptionSchema, ChoiceOptionSchema


class BaseLLMService:
    _llm = None
    _instance = None

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            new_instance = super(BaseLLMService, cls).__new__(cls)
            new_instance._llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
            cls._instance = new_instance
        return cls._instance


class RadioMatchService(BaseLLMService):
    """单项选择匹配服务"""

    template = """您将收到一个单项选择题。您的目标时根据题目找出选项。您需要注意有符合选项时只返回选项内容。您需要注意没有符合的选项时返回空。
{format_instructions}

选项以(``)包裹。
`{params}`

题目将以(```)包裹.
```{query}```
最终答案："""
    template_variables = ["query", "params"]
    output_key = "result"
    _empty_result = "空"

    def __init__(self):
        self._parser = PydanticOutputParser(pydantic_object=OptionSchema)
        prompt = PromptTemplate(
            template=self.template,
            input_variables=self.template_variables,
            partial_variables={"format_instructions": self._parser.get_format_instructions()}
        )
        # 单项选择搜索链
        self._chain = LLMChain(
            llm=self._llm,
            prompt=prompt,
            output_key=self.output_key
        )

        super(BaseLLMService, self).__init__()

    def _format_result(self, result) -> OptionSchema:
        """
        匹配返回结果
        :param result:
        :return:
        """
        json_str = result.get(self.output_key)
        if not result or result == self._empty_result:
            return OptionSchema()

        option_parser = self._parser

        try:
            res = option_parser.parse(json_str)
        except Exception as e:
            try:
                llm = self._llm
                new_parser = OutputFixingParser.from_llm(parser=option_parser, llm=llm)
                res = new_parser.parse(json_str)
            except Exception as e:
                print(json_str, e)
                res = None
        return res

    async def search(self, question: str, options: str) -> str:
        """
        从属性中匹配关键词
        :param question: 问题
        :param options: 选项
        :return: 符合的选项
        """
        result = await self._chain.ainvoke(dict(query=question, params=options))
        option = self._format_result(result)
        return option.option


class ChoiceMatchService(BaseLLMService):
    """多项选择题匹配"""

    template = """您将收到一个多项选择题。您的目标时根据题目找出选项。您需要根据题目和选项的语义，尽可能多的匹配语义相近选项。您需要注意没有符合的选项时返回空。
{format_instructions}

选项以(``)包裹。
`{params}`

题目将以(```)包裹.
```{query}```  

符合的选项："""

    template_variables = ["query", "params"]
    output_key = "result"
    _empty_result = "空"

    def __init__(self):
        self._parser = PydanticOutputParser(pydantic_object=ChoiceOptionSchema)
        prompt = PromptTemplate(
            template=self.template,
            input_variables=self.template_variables,
            partial_variables={"format_instructions": self._parser.get_format_instructions()}
        )
        # 单项选择搜索链
        self._chain = LLMChain(
            llm=self._llm,
            prompt=prompt,
            output_key=self.output_key
        )

        super(BaseLLMService, self).__init__()

    def _format_result(self, result) -> ChoiceOptionSchema:
        """
        匹配返回结果
        :param result:
        :return:
        """
        res = ChoiceOptionSchema()
        json_str = result.get(self.output_key)
        if not result or result == self._empty_result:
            return ChoiceOptionSchema()

        parser = self._parser
        llm = self._llm

        try:
            res = parser.parse(json_str)
        except Exception as e:
            try:
                new_parser = OutputFixingParser.from_llm(parser=parser, llm=llm)
                res = new_parser.parse(json_str)
            except Exception as e:
                print(json_str, e)
        return res

    async def search(self, question: str, options: str) -> List:
        """
        从属性中匹配关键词
        :param question: 问题
        :param options: 选项
        :return: 符合的选项
        """
        result = await self._chain.ainvoke(dict(query=question, params=options))
        choice = self._format_result(result)
        return [x.option for x in choice.lists] if choice else []


class MatchMeterService(BaseLLMService):
    """
    匹配设备属性列表
    """
    template = """您是一个音视频设备专家。您将收到一个设备的属性。您的目标是从设备属性中找出目标参数：`{params}`。您需要按照目标参数返回。
{format_instructions}
设备属性将以(```)包裹.
```{query}```
最终答案："""

    template_variables = ["query", "params"]
    output_key = "result"
    _empty_result = "空"

    def __init__(self):
        self._parser = PydanticOutputParser(pydantic_object=DeviceAttributeListSchema)
        prompt = PromptTemplate(
            template=self.template,
            input_variables=self.template_variables,
            partial_variables={"format_instructions": self._parser.get_format_instructions()}
        )

        self._chain = LLMChain(
            llm=self._llm,
            prompt=prompt,
            output_key=self.output_key
        )
        super(BaseLLMService, self).__init__()

    def _format_result(self, result) -> DeviceAttributeListSchema:
        """
        匹配返回结果
        :param result:
        :return:
        """
        json_str = result.get(self.output_key)
        if not result or result == self._empty_result:
            return DeviceAttributeListSchema()

        option_parser = self._parser

        try:
            res = option_parser.parse(json_str)
        except Exception as e:
            try:
                llm = self._llm
                new_parser = OutputFixingParser.from_llm(parser=option_parser, llm=llm)
                res = new_parser.parse(json_str)
            except Exception as e:
                print(json_str, e)
                res = None
        return res

    async def search(self, meter: str, keywords: str) -> List:
        """
        从属性中找到关键设备信息
        :param meter: 设备属性
        :param keywords: 关键属性
        :return: [DeviceAttributeSchema] 匹配到的属性列表
        """
        result = await self._chain.ainvoke(dict(query=meter, params=keywords))
        option = self._format_result(result)
        return option.lists if option else []

